package bjut.zrpc.core.optimization.serializer;

import bjut.zrpc.core.optimization.interfaces.Serializer;
import bjut.zrpc.core.optimization.message.ZRpcRequestMessage;
import bjut.zrpc.core.optimization.message.ZRpcResponseMessage;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;


import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * @Author: zhuzhiming
 * @Date: 2022/6/24 13:04
 */
public class KryoSerializer implements Serializer {

    private static final ThreadLocal<Kryo> kryoThreadLocal = ThreadLocal.withInitial(() -> {
        Kryo kryo = new Kryo();
        kryo.register(ZRpcResponseMessage.class);
        kryo.register(ZRpcRequestMessage.class);
        return kryo;
    });

    @Override
    public byte[] serialize(Object data) {

        try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream()) {

            Output output = new Output(byteArrayOutputStream);
            Kryo kryo = kryoThreadLocal.get();
            kryo.writeObject(output, data);
            kryoThreadLocal.remove();
            return output.toBytes();

        } catch (IOException e) {
            throw new RuntimeException(e.getMessage());
        }

    }

    @Override
    public <T> T deSerialize(byte[] bytes, Class<T> cls) {

        try (ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes)) {

            Input input = new Input(byteArrayInputStream);
            Kryo kryo = kryoThreadLocal.get();
            T t = kryo.readObject(input, cls);
            return cls.cast(t);

        } catch (IOException e) {
            throw new RuntimeException(e.getMessage());
        }

    }
}
